/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package sk.smitala.gephi.communitydetection.utils;

import java.util.HashMap;
import org.gephi.data.attributes.api.AttributeColumn;
import org.gephi.graph.api.Node;
import sk.smitala.gephi.communitydetection.utils.traverse.TraverseListener;

/**
 *
 * @author smitalm
 */
public class GraphPartitionLoader implements TraverseListener<Node> {

    AttributeColumn attribute;
    HashMap<String, Community> communities = new HashMap<String, Community>();

    public GraphPartitionLoader(AttributeColumn attribute) {
	this.attribute = attribute;
    }

    @Override
    public void onElementVisited(Node node) {
	String attributeID = attribute.getId();
	String value;
	Object obj = node.getNodeData().getAttributes().getValue(attributeID);
	if (obj instanceof String) {
	    value = (String) obj;
	} else if (obj instanceof Integer) {
	    value = ((Integer) obj).intValue() + "";
	} else {
	    return;
	}
	if (communities.containsKey(value)) {
	    communities.get(value).addMember(node.getId());
	} else {
	    communities.put(value, new Community(value, node.getId()));
	}
    }

    public HashMap<String, Community> getPartition() {
	return communities;
    }
}
